package com.os.opencv.java.chapter12;

import org.opencv.core.*;
import org.opencv.highgui.HighGui;
import org.opencv.imgproc.Imgproc;
import org.opencv.ml.Ml;
import org.opencv.ml.SVM;

public class SVM1 {

    public static void main(String[] args) {
        System.loadLibrary(Core.NATIVE_LIBRARY_NAME);

        //训练数据
        float[] trainData = {550,50,300,50,500,250,50,550};
        Mat trainMat = new Mat(4,2, CvType.CV_32FC1);
        trainMat.put(0,0,trainData);

        //训练数据标签
        int[] label = {1,-1,-1,-1};
        Mat labelMat = new Mat(4,1,CvType.CV_32SC1);
        labelMat.put(0,0,label);

        //创建svm对象并设置参数
        SVM svm = SVM.create();
        svm.setType(SVM.C_SVC);
        svm.setKernel(SVM.LINEAR);
        svm.setTermCriteria(new TermCriteria(TermCriteria.MAX_ITER, 100, 1e-6));

        //训练svm分类器
        svm.train(trainMat, Ml.ROW_SAMPLE, labelMat);

        //创建显示测试结果图像dst
        int width = 600;
        int height = 600;
        Mat dst = Mat.zeros(height, width, CvType.CV_8UC3);
        int col = dst.cols();
        int c = dst.channels();

        //测试每个像素并将结果存入结果图像
        byte[] dstData = new byte[(int)(dst.total() * dst.channels())];

        Mat testMat = new Mat(1,2,CvType.CV_32FC1);
        float[] testData = new float[(int)2];
        for(int i=0; i<dst.rows(); i++){
            for(int j=0; j<dst.cols(); j++){
                //对每个像素值进行测试
                testData[0] = j;
                testData[1] = i;
                testMat.put(0,0,testData);
                float result = svm.predict(testMat);

                //根据测试结果用不同颜色标记
                if(result == 1){
                    dstData[(i*col+j) * c] = 0;
                    dstData[(i*col+j) * c + 1] = (byte) 255;
                    dstData[(i*col+j) * c + 2] = 0;
                }else if(result == -1){
                    dstData[(i*col+j) * c] = (byte) 255;
                    dstData[(i*col+j) * c + 1] = 0;
                    dstData[(i*col+j) * c + 2] = 0;
                }
            }
        }

        //将所有数据放入结果图像中
        dst.put(0,0,dstData);

        //用黑色和白色圆圈画出训练数据
        Scalar white = new Scalar(255,255,255);
        Scalar black = new Scalar(0,0,0);
        Scalar red = new Scalar(0,0,255);
        Imgproc.circle(dst, new Point(trainData[0], trainData[1]), 5, black, -1);
        Imgproc.circle(dst, new Point(trainData[2], trainData[3]), 5, black, -1);
        Imgproc.circle(dst, new Point(trainData[4], trainData[5]), 5, black, -1);
        Imgproc.circle(dst, new Point(trainData[6], trainData[7]), 5, black, -1);

        //获取支持向量
        Mat sv = svm.getUncompressedSupportVectors();
        float[] svData = new float[(int)(sv.total() * sv.channels())];
        sv.get(0,0,svData);
        //在结果图像上有矩形画出支持向量
        for(int i=0; i<sv.rows(); i++){
            double x = svData[i*sv.cols()];
            double y = svData[i*sv.cols() + 1];
            Imgproc.rectangle(dst, new Point(x-8, y-8), new Point(x+8, y+8), red, 2);
        }

        //在屏幕上显示结果图像
        HighGui.imshow("svm", dst);
        HighGui.waitKey();
        System.exit(0);

    }
}
